'''
作为预测代码，可用于模型的预测
1. tripletloss_3D_train.py: tripletloss_3D_*
'''

import os
import pandas as pd
import numpy as np

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from torchvision.models import resnet50

from sklearn.metrics import classification_report, confusion_matrix
from net_tripletloss_3d import generate_model

BatchSize = 16
save_path = '/home/zmy/pytorch_code/checkpoint/tripletloss_3D_0.1_20210509_114838'
weight_name = 'tripletloss_3D_0.1-80-3.3748047153155007-regular.pth'

weights_path = save_path + '/' + weight_name


# 自定义数据集和数据预处理
class MyDataset(Dataset):

    def __init__(self, datalist):
        self.data_info = datalist

    def __len__(self):
        return len(self.data_info)

    def __getitem__(self, item):
        patientID = self.data_info[item][0]
        label = self.data_info[item][1]
        img_path = self.data_info[item][2]

        img = np.load(img_path)

        # name_list = os.listdir(ct_path)
        # name_list.sort()
        #
        # ct_3D_array = []
        # pet_3D_array = []
        #
        # for i in range(len(name_list)):
        #     ct = np.load(ct_path+name_list[i])
        #     pet = np.load(pet_path+name_list[i])
        #
        #     ct_3D_array.append(ct)
        #     pet_3D_array.append(pet)
        #
        #
        # # ct和pet进行合并
        # img = np.asarray([ct_3D_array, pet_3D_array], dtype=np.float)
        #
        # # print('img shape: ', img.shape)
        #
        # # 归一化输入大小
        # resize_para = 8.0 / len(name_list)
        # img = scipy.ndimage.zoom(img, [1, resize_para, 0.5, 0.5])
        #
        # # print('resize img shape: ', img.shape)

        return {'image': torch.from_numpy(img), 'label': torch.tensor(label)}


# 读取文件列表
def read_csv(data_sets):
    sets_path = '/data1/zmy/data2021/origin_data/divide_csv/five/'

    # 读取数据集
    data_features = []

    for set in data_sets:
        train_data = pd.read_csv(sets_path + set)
        for j in range(len(train_data)):
            # 读取文件地址
            patientid = train_data['patientID'][j]

            img_path = '/data1/zmy/data2021/auge2_data/img_3D/' + str(patientid) + '/img3d.npy'

            one_feature = [patientid, int(train_data['cancer_type'][j]) - 1, img_path]

            data_features.append(one_feature)

    return data_features


# 创建网络
def resnet():
    # 创建resnet50网络
    net = generate_model(50)

    # 打印网络结构和参数量
    # print(net)
    # print("Total number of paramerters in networks is {}  ".format(sum(x.numel() for x in net.parameters())))

    return net


# 结果预测
def predict():
    # 创建网络结构
    net = resnet().to(device)

    # 加载测试数据
    test_data_sets = ['test.csv']
    test_list = read_csv(test_data_sets)
    test_dataset = MyDataset(test_list)
    testloader = DataLoader(test_dataset, batch_size=BatchSize, shuffle=False, num_workers=2)

    # 加载模型权重
    net.load_state_dict(torch.load(weights_path))
    net.eval()

    # 返回的预测结果及对应标签
    y_preds = []
    y_true = []

    with torch.no_grad():
        for n_iter, data in enumerate(testloader):
            print('iteration:{}\ttotal {} iterations'.format(n_iter + 1, len(testloader)))

            images = data['image'].type(torch.FloatTensor).to(device)
            labels = data['label'].to(device)

            out1, outputs = net(images)
            _, preds = outputs.max(1)

            y_preds.extend(preds.tolist())
            y_true.extend(labels.tolist())

            print(preds.tolist())
            print(labels.tolist())
    return y_preds, y_true


# 结果评估
def evaluation(preds, labels):
    # 计算评估指标
    target_names = ['1', '2', '3', '4', '5']
    result_statis = classification_report(y_true=labels, y_pred=preds, target_names=target_names)
    print(result_statis)

    # 计算混淆矩阵
    confusion = confusion_matrix(y_true=labels, y_pred=preds)
    print(confusion)


if __name__ == '__main__':
    device = torch.device("cuda:3")

    # 预测结果
    y_preds, y_true = predict()

    print('weight name: ', weight_name)

    # 结果评估
    evaluation(y_preds, y_true)


